"""核心层：ProofMapper 会话态管理（基于 UUID 的 shared 管理）

目标：
- 统一 final_output 为 dict（已在节点层修正：`nodes/interactive_nodes.py`）。
- 核心层维护以 UUID 标识的 shared，前端仅需持有 shared_id 进行后续交互/回滚。

对外提供：
- create_session_async/pdf 同步封装：从 PDF 解析，创建会话，返回 shared_id。
- get_shared：获取当前会话的精简 shared（final_output/modification_log/error_log）。
- apply_user_input_async/_by_id：基于自然语言修改指定会话。
- rollback_shared_by_id：回滚指定会话。
- delete_session：删除会话。

实现细节：
- 内存级单例会话存储，满足单用户/开发阶段需求。
- 若需多实例/重启持久化，可替换为 Redis 等外部存储。
"""

import threading
import uuid
from datetime import datetime
from typing import Any, Dict, Optional, Tuple
from log import log_info

# 复用接口层的功能实现
from agent.ProofMapper.api import (
    parse_pdf_to_shared_async,
    apply_user_input_async,
    rollback_shared,
)


# -----------------------------
# 内存会话存储（单例）
# -----------------------------
class _SharedStore:
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._store: Dict[str, Dict[str, Any]] = {}

    def set(self, key: str, value: Dict[str, Any]) -> None:
        with self._lock:
            self._store[key] = value

    def get(self, key: str) -> Optional[Dict[str, Any]]:
        with self._lock:
            return self._store.get(key)

    def delete(self, key: str) -> bool:
        with self._lock:
            return self._store.pop(key, None) is not None


_shared_store = _SharedStore()


# -----------------------------
# 工具方法
# -----------------------------
def _ensure_simplified_shared(shared: Dict[str, Any]) -> Dict[str, Any]:
    """确保 shared 仅包含精简键并补齐默认值。"""
    return {
        "final_output": shared.get("final_output", {}) or {},
        "modification_log": shared.get("modification_log", []) or [],
        "error_log": shared.get("error_log", []) or [],
    }


# -----------------------------
# 会话创建（解析 PDF）
# -----------------------------
async def create_session_async(pdf_path: str, output_path: str = "output.json") -> Tuple[str, Dict[str, Any]]:
    """解析 PDF，创建会话并返回 (shared_id, simplified_shared)。"""
    log_info("========开始解析PDF========","proofmapper")
    shared = await parse_pdf_to_shared_async(pdf_path, output_path)
    simplified = _ensure_simplified_shared(shared)
    shared_id = uuid.uuid4().hex
    _shared_store.set(shared_id, simplified)
    return shared_id, simplified


# -----------------------------
# 从快照/现成 JSON 创建会话（仅异步接口）
# -----------------------------
async def create_session_from_json_async(data: Dict[str, Any]) -> Tuple[str, Dict[str, Any]]:
    """从快照或最终JSON创建会话，返回 (shared_id, simplified_shared)。

    兼容两种输入：
    - 最终JSON结构：包含 'items'（以及可选 'metadata'）
    - 快照结构：包含 'evaluation_items'、'assessment_type'、'name'
    """
    final_json: Dict[str, Any] = {}

    if isinstance(data, dict):
        if isinstance(data.get("items"), list):
            # 已是最终JSON
            final_json = data
        elif isinstance(data.get("evaluation_items"), list):
            # 从快照结构转换为最终JSON
            final_json = {
                "metadata": {
                    "source_file": data.get("name") or data.get("assessment_type") or "snapshot",
                    "generated_at": datetime.now().isoformat(),
                    "version": "1.0",
                },
                "items": data.get("evaluation_items", []),
            }

    simplified = {
        "final_output": final_json if final_json else {},
        "modification_log": [],
        "error_log": [] if final_json else ["invalid snapshot json"],
    }

    shared_id = uuid.uuid4().hex
    _shared_store.set(shared_id, simplified)
    return shared_id, simplified




# -----------------------------
# 会话读取（仅异步接口）
# -----------------------------
async def get_shared_async(shared_id: str) -> Optional[Dict[str, Any]]:
    """获取指定会话的精简 shared（异步）。"""
    shared = _shared_store.get(shared_id)
    if shared is None:
        return None
    return _ensure_simplified_shared(shared)



# -----------------------------
# 会话内交互修改
# -----------------------------
async def apply_user_input_async_by_id(shared_id: str, user_input: str) -> Optional[Dict[str, Any]]:
    """对指定会话应用一次自然语言修改，返回更新后的精简 shared。"""
    shared = _shared_store.get(shared_id)
    if shared is None:
        return None

    log_info("========开始自然语言交互修改========","proofmapper")
    updated = await apply_user_input_async(user_input, shared)
    simplified = _ensure_simplified_shared(updated)

    # 将 message 追加到 error_log（无论成功或失败），便于前端统一读取提示信息
    msg = updated.get("message")
    if msg:
        simplified.setdefault("error_log", []).append(msg)

    _shared_store.set(shared_id, simplified)
    return simplified



# -----------------------------
# 会话回滚（仅异步接口）
# -----------------------------
async def rollback_shared_async_by_id(shared_id: str) -> Optional[Dict[str, Any]]:
    """对指定会话回滚最近一次修改，返回更新后的精简 shared（异步）。"""
    shared = _shared_store.get(shared_id)
    if shared is None:
        return None

    log_info("========开始回滚========","proofmapper")
    updated = rollback_shared(shared)
    simplified = _ensure_simplified_shared(updated)
    _shared_store.set(shared_id, simplified)
    return simplified


# -----------------------------
# 删除会话（仅异步接口）
# -----------------------------
async def delete_session_async(shared_id: str) -> bool:
    """删除指定会话，释放内存（异步）。"""
    return _shared_store.delete(shared_id)



